#!/usr/bin/env python
# -*- coding: utf-8 -*-

# http://docs.cython.org/docs/tutorial.html#cython-hello-world
# http://nedbatchelder.com/text/whirlext.html
# import pyximport; pyximport.install(); import tctable


'''This Cython module contains the TCTable class, a driver for
the table database in Tokyo Cabinet.

The TCQuery class represents a query to the table database.

TO DO
=====

* TCTable.get() should accept a *cols* argument to make it possible for the
user to specify the desired columns.

* Implement the missing functions (e.g. tctdbsetindex(), transactions,
performance config).
'''


cdef extern from "Python.h":
    ctypedef struct PyObject
    int PyString_AsStringAndSize(object obj, char **buffer, int *length) except -1
    #ctypedef struct PyBytesObject
    #object PyString_FromStringAndSize(char *s, Py_ssize_t len)
    object PyBytes_FromStringAndSize(char *, Py_ssize_t)
cdef extern from "tctdb.h":
    char *c_tcversion "tcversion"
    ctypedef unsigned char uint8_t
    ctypedef unsigned char  int8_t
    ctypedef int           int32_t
    ctypedef int          uint32_t
    ctypedef int           int64_t
    ctypedef int          uint64_t
    ctypedef int           _Bool
    ctypedef struct TDBIDX:
        pass
    ctypedef struct TCTDB:
        void *mmtx
        #TCHDB *hdb
        _Bool open
        _Bool wmode
        uint8_t opts
        int32_t lcnum
        int32_t ncnum
        #TDBIDX *idxs
        int   inum
        _Bool tran
    ctypedef struct TDBQRY:
        pass
    cdef enum:                 # Query conditions
        TDBQCSTREQ             # string is equal to */
        TDBQCSTRINC            # string is included in */
        TDBQCSTRBW             # string begins with */
        TDBQCSTREW             # string ends with */
        TDBQCSTRAND            # string includes all tokens in */
        TDBQCSTROR             # string includes at least one token in */
        TDBQCSTROREQ           # string is equal to at least one token in */
        TDBQCSTRRX             # string matches regular expressions of */
        TDBQCNUMEQ             # number is equal to */
        TDBQCNUMGT             # number is greater than */
        TDBQCNUMGE             # number is greater than or equal to */
        TDBQCNUMLT             # number is less than */
        TDBQCNUMLE             # number is less than or equal to */
        TDBQCNUMBT             # number is between two tokens of */
        TDBQCNUMOREQ           # number is equal to at least one token in */
        TDBQCNEGATE = 1 << 24  # negation flag */
        TDBQCNOIDX = 1 << 25   # no index flag */
    cdef enum:                 # Query order
        TDBQOSTRASC
        TDBQOSTRDESC
        TDBQONUMASC
        TDBQONUMDESC
    ctypedef struct TCLISTDATUM: # type of structure for an element of a list */
        char *ptr            # pointer to the region */
        int size      # size of the effective region */
    ctypedef struct TCLIST: # type of structure for an array list */
        TCLISTDATUM *array  # array of data */
        int anum            # number of the elements of the array */
        int start           # start index of used elements */
        int num             # number of used elements */
    ctypedef struct TCMAPREC:       # an element of a map
        int ksiz          # size of the region of the key
        int vsiz          # size of the region of the value
        unsigned int hash # second hash value
        TCMAPREC *left    # pointer to the left child
        TCMAPREC *right   # pointer to the right child
        TCMAPREC *prev    # pointer to the previous element
        TCMAPREC *next
    ctypedef struct TCMAP: # represents a table row
        TCMAPREC **buckets # bucket array
        TCMAPREC *first    # pointer to the first element
        TCMAPREC *last     # pointer to the last element
        TCMAPREC *cur      # pointer to the current element
        uint32_t bnum      # number of buckets
        uint64_t rnum      # number of records
        uint64_t msiz      # total size of records
    int tctdbstrtoindextype(char *str)
    char *tctdberrmsg      (int ecode)
    TCTDB *tctdbnew    ()
    void tctdbdel      (TCTDB *tdb)
    int tctdbecode     (TCTDB *tdb)
    _Bool tctdbsetmutex(TCTDB *tdb)
    _Bool tctdbtune(TCTDB *tdb, int64_t bnum, int8_t apow, int8_t fpow, uint8_t opts)
    _Bool tctdbsetcache(TCTDB *tdb, int32_t rcnum, int32_t lcnum, int32_t ncnum)
    _Bool tctdbsetxmsiz(TCTDB *tdb, int64_t xmsiz)
    _Bool tctdbopen    (TCTDB *tdb, char *path, int omode)
    _Bool tctdbclose   (TCTDB *tdb)
    _Bool tctdbsync    (TCTDB *tdb) # flush
    _Bool  tctdbput    (TCTDB *tdb, void *pkbuf, int pksiz, TCMAP *cols)
    _Bool  tctdbout    (TCTDB *tdb, void *pkbuf, int pksiz) # deletes a row
    int   tctdbvsiz    (TCTDB *tdb, void *pkbuf, int pksiz) # row size
    uint64_t tctdbrnum (TCTDB *tdb)    # number of records in table
    _Bool tctdbputkeep (TCTDB *tdb, void *pkbuf, int pksiz, TCMAP *cols)
    _Bool tctdbiterinit(TCTDB *tdb)
    # Maps
    TCMAP *tctdbget    (TCTDB *tdb, void *pkbuf, int pksiz)
    TCMAP *tcmapnew2(uint32_t bnum) # Create a map specifying number of buckets
    void tcmapdel      (TCMAP *map)
    void tcmapput      (TCMAP *map, void *kbuf, int ksiz, void *vbuf, int vsiz)
    void tcmapiterinit (TCMAP *map)
    void *tcmapiternext(TCMAP *map, int *sp)
    void *tcmapget     (TCMAP *map, void *kbuf, int ksiz, int *sp)
    # Queries
    TDBQRY *tctdbqrynew    (TCTDB *tdb) # create query
    void tctdbqrydel       (TDBQRY *qry)   # release query
    TCLIST *tctdbqrysearch (TDBQRY *qry)   # execute
    _Bool tctdbqrysearchout(TDBQRY *qry)   # execute deleting all matches
    void tctdbqryaddcond   (TDBQRY *qry, char *name, int op, char *expr)
    void tctdbqrysetorder  (TDBQRY *qry, char *name, int type)
    void tctdbqrysetlimit  (TDBQRY *qry, int max, int skip)
    void tclistdel  (TCLIST *list)
    int  tclistnum  (TCLIST *list)
    void *tclistval (TCLIST *list, int index, int *sp)
    char *tclistval2(TCLIST *list, int index)


tokyo_cabinet_version = c_tcversion
__version__ = '0.1'


class TCException(Exception):
    '''Represents an error communicating with the Tokyo Cabinet database.'''

    

cdef class TCTable(object):
    '''A Tokyo Cabinet table database.
    
    The TCTable constructor takes these parameters:
    
    *path* is the path to the database file to be opened. Its extension
    is normally .tct.
    
    *mode* is either 'w' for write or 'r' for read.
    
    The object is used exactly as a Python dictionary; however,
    all its values are dicts containing column:value pairs, where both
    *column* and *value* must be instances of *bytes*.
    
    TODO: TCTable is not iterable; currently you have to tctable.query.all()
    
    TODO: Add transactions, vanish, tctdbcopy (backup).
    '''
    cdef TCTDB* _state
    cdef readonly char *path
    cdef readonly char *mode
    cdef _Bool is_open
    
    def _throw(self, message='Error.'):
        '''Raises a TCException, appending to *message* the last error message
        from the table database.
        '''
        raise TCException(self._msg(message))
    
    def _msg(self, message='Error.'):
        '''Composes an error message for an exception, appending to the value
        of *message* the last error message from the table database.
        '''
        cdef int errorcode = tctdbecode(self._state)
        cdef char *msg = <char*>tctdberrmsg(errorcode)
        return message + ' ' + msg.capitalize()
    
    def __init__(self, path, mode='r'):
        self.path = path
        self.mode = mode
        self._state = tctdbnew()
        cdef _Bool success = tctdbopen(self._state, path, 6 if mode=='w' else 1)
        if not success:
            self._throw('Unable to open {0} for {1}.'.format \
                (str(path), 'writing' if mode == 'w' else 'reading'))
        self.is_open = True
    
    def close(self):
        '''Closes the database file and cleans memory, rendering this
        object useless.
        '''
        #print 'Closing ' + str(self.path)
        #cdef _Bool success = tctdbclose(self._state)
        if self.is_open:
            self.is_open = False
            tctdbdel(self._state)
    
    def __setitem__(self, k, dic):
        '''Takes a key and a value and writes them as a table row.
        The key must be a bytes instance.
        The value must be a dictionary containing bytes.
        '''
        # When you alter this method, you probably modify keep_or_put() too.
        # Copy the dictionary to a TCMAP
        cdef TCMAP * tcmap = tcmapnew2(len(dic))
        cdef char *kbuf, *vbuf
        cdef int ksize, vsize
        for key, val in dic.items():
            # http://bitbucket.org/james_taylor/bx-python/src/tip/lib/bx/_seqmapping.pyx
            PyString_AsStringAndSize(key, &kbuf, &ksize)
            PyString_AsStringAndSize(val, &vbuf, &vsize)
            tcmapput(tcmap, kbuf, ksize, vbuf, vsize)
        # Put the key in C variable
        PyString_AsStringAndSize(k, &kbuf, &ksize)
        # Write the TCMAP into the DB
        cdef _Bool success = tctdbput(self._state, kbuf, ksize, tcmap)
        tcmapdel(tcmap)
        if not success:
            self._throw('Unable to write to database '+ str(k))
    
    def keep_or_put(self, k, dic):
        '''Takes as arguments a key (string) and a value (dict).
        If the key already exists in the table, nothing is done
        and the string 'keep' is returned.
        Otherwise, the value is stored at the key and 'put' is returned.
        '''
        # When you alter this method, you probably modify __setitem__() too.
        # Copy the dictionary to a TCMAP
        cdef TCMAP * tcmap = tcmapnew2(len(dic))
        cdef char *kbuf, *vbuf
        cdef int ksize, vsize
        for key, val in dic.items():
            # http://bitbucket.org/james_taylor/bx-python/src/tip/lib/bx/_seqmapping.pyx
            PyString_AsStringAndSize(key, &kbuf, &ksize)
            PyString_AsStringAndSize(val, &vbuf, &vsize)
            tcmapput(tcmap, kbuf, ksize, vbuf, vsize)
        # Put the key in C variable
        PyString_AsStringAndSize(k, &kbuf, &ksize)
        # Keep existing, or write the TCMAP into the DB
        cdef _Bool success = tctdbputkeep(self._state, kbuf, ksize, tcmap)
        tcmapdel(tcmap)
        cdef int errorcode
        if not success:
            errorcode = tctdbecode(self._state)
            if errorcode == 21: # 'existing record'
                return 'keep'
            else:
                self._throw('Unable to write key "{0}".'.format(str(k)))
        return 'put'
    
    def setdefault(self, k, dic):
        '''Writes a record -- d[k] = dic -- if k not in d.
        Then returns d[k].
        '''
        if self.keep_or_put(k, dic) == 'keep':
            return self.get(k, dic)
        else:
            return dic # which certainly has just been put
    
    def __getitem__(self, id):
        '''Returns a record (as a dict) or raises KeyError.'''
        # If you alter this method, you probably want to alter get() too.
        cdef char *buffer
        cdef int   buffersize
        # Put the row key in C variable
        PyString_AsStringAndSize(id, &buffer, &buffersize)
        # Retrieve the tcmap
        cdef TCMAP *tcmap = tctdbget(self._state, buffer, buffersize)
        if tcmap == NULL:
            raise KeyError('Lookup failed: ' + str(id))
        #print 'Size:',tcmap.msiz, 'Cols:',tcmap.rnum, 'Buckets:',tcmap.bnum
        # Translate tcmap to a Python dict
        dic = {}
        tcmapiterinit(tcmap) # Initialize the map iterator
        cdef void *key, *val
        cdef int keysize, valsize
        while True:
            key = tcmapiternext(tcmap, &keysize) # Get one column
            if key == NULL: break
            val = tcmapget     (tcmap, key, keysize, &valsize)
            #print keysize, <char*>key, valsize, <char*>val
            dic[PyBytes_FromStringAndSize(<char*>key, keysize)] = \
                PyBytes_FromStringAndSize(<char*>val, valsize)
        tcmapdel(tcmap)
        return dic
    
    def get(self, id, default=None):
        '''Returns a record (as a dict) or the value of the keyword argument
        *default*.
        '''
        # If you alter this method, modify __getitem__() too.
        cdef char *buffer
        cdef int   buffersize
        # Put the row key in C variable
        PyString_AsStringAndSize(id, &buffer, &buffersize)
        # Retrieve the tcmap
        cdef TCMAP *tcmap = tctdbget(self._state, buffer, buffersize)
        if tcmap == NULL:
            return default
        #print 'Size:',tcmap.msiz, 'Cols:',tcmap.rnum, 'Buckets:',tcmap.bnum
        # Translate tcmap to a Python dict
        dic = {}
        tcmapiterinit(tcmap) # Initialize the map iterator
        cdef void *key, *val
        cdef int keysize, valsize
        while True:
            key = tcmapiternext(tcmap, &keysize) # Get one column
            if key == NULL: break
            val = tcmapget     (tcmap, key, keysize, &valsize)
            #print keysize, <char*>key, valsize, <char*>val
            dic[PyBytes_FromStringAndSize(<char*>key, keysize)] = \
                PyBytes_FromStringAndSize(<char*>val, valsize)
        tcmapdel(tcmap)
        return dic
    
    def __delitem__(self, id):
        cdef char *buffer
        cdef int   buffersize
        # Put the key in C variable
        PyString_AsStringAndSize(id, &buffer, &buffersize)
        # Delete the row
        cdef _Bool success = tctdbout(self._state, buffer, buffersize)
        if not success:
            raise KeyError(self._msg('Could not delete row "{0}".' \
                                     .format(str(id))))
    
    def pop(self, k, d=None):
        '''D.pop(k[,d]) -> v, remove specified key and return the
        corresponding value. If key is not found, d is returned if given,
        otherwise KeyError is raised.
        '''
        adict = self.get(k)
        if adict is None:
            if d is None:
                raise KeyError('Key not found: "{0}"'.format(k))
            else:
                return d
        else:
            del self.k
            return adict
    
    def size_of(self, id):
        '''Returns the size of the value stored at the passed *id*,
        or -1 if the id does not exist in the table database.
        
        I don't know why this would be useful... the returned number is
        a little higher than the amount of characters persisted.
        '''
        cdef char *buffer
        cdef int   buffersize
        # Put the row key in C variable
        PyString_AsStringAndSize(id, &buffer, &buffersize)
        # Retrieve the size and return it
        cdef int size = tctdbvsiz(self._state, buffer, buffersize)
        return size
    
    def __contains__(self, key):
        '''True if D has a key *key*, else False.
        But know that calling D.get(id, defaultdict) is cheaper than
            dic = D[id] if id in D else defaultdict
        '''
        return self.size_of(key) > -1
    
    def flush(self):
        '''Forces changes to be written to disk. Returns a success code.'''
        cdef _Bool success = tctdbsync(self._state)
        return success
    
    def __len__(self):
        cdef uint64_t number = tctdbrnum(self._state)
        return number
    
    @property
    def query(self):
        '''Returns a new query object for searching this table.'''
        return make_query(self, self._state)
    
    """def __iter__(self):
        '''Gets an iterator that yields all keys in the TCTable.'''
        cdef TCTDB* _state = self._state
        i = TableIter(_state)
        return i
    """


cdef class TCQuery(object):
    '''Represents a query to the table database. You should never instantiate
    this yourself. Instead, obtain an instance from the query property of
    your table:
    
        # Query for all objects whose type is NOT 'person'
        query = mytable.query.col('type').no.equals('person')
        all_IDs = query.keys() # a list
        print "Deleting ", all_IDs
        query.delete_all()
    '''
    cdef TDBQRY* _state
    cdef TCLIST *result
    cdef int _no
    cdef str _col
    cdef object _table
    ''' Not sure when or whether to call tctdbqrydel().
    def close(self):
        tctdbqrydel(self._state)
    # If __del__() is defined, the instance cannot be collected by the GC.
    def __del__(self):
        tctdbqrydel(self._state)
    '''
    
    def count(self):
        '''Executes the query and returns the number of matching elements.
        Only use this when you really care only for the number of rows.
        If you call count() to know the number of rows, then keys() to get them
        you are actually performing the query twice, so don't.
        '''
        cdef TCLIST *tclist = tctdbqrysearch(self._state)
        cdef int count = tclistnum(tclist) # number of elements in the list
        return count
    
    def keys(self):
        '''Executes the query and returns a list of the matching IDs...
        or an empty list.
        '''
        cdef TCLIST *tclist = tctdbqrysearch(self._state)
        cdef int count = tclistnum(tclist) # number of elements in the list
        pylist = []
        ''' I wish this version of the for loop would work.
        cdef int *areasize
        cdef void *area
        for i in range(count):
            print 'Iteration', i
            area = tclistval(tclist, i, areasize) # get area of the value
            if area == NULL:
                print 'NULL !!!!!111 BUG'
                break
            print 'got area', <char*>area
            b = PyBytes_FromStringAndSize(<char*>area, areasize[0])
            print b
            pylist.append(b)
        '''
        cdef char *chars
        for i in range(count):
            chars = tclistval2(tclist, i)
            pylist.append(chars)
        tclistdel(tclist)
        #self.close()
        return pylist
    
    def delete_all(self):
        '''Executes the query and removes all corresponding objects.'''
        cdef _Bool success = tctdbqrysearchout(self._state)
        if not success:
            raise TCException('Unable to delete all objects in query.')
        return self
    
    def all(self):
        '''Executes the query and returns a list of tuples (key, dict)
        containing the matching rows.
        '''
        keys = self.keys()
        alist = []
        for key in keys:
            alist.append((key, self._table.get(key)))
        return alist
    
    def col(self, name):
        '''Sets the column to which will be applied the next condition
        called on this query.
        '''
        self._col = name
        return self
    
    def order_by(self, col, desc=False, number=False):
        '''Sets the order of the query results, with these arguments:
        *col*: the name of the column to order by
        *desc*: descending order. The default is False (ascending).
        *number*: use the number sorting algorithm. Default is False
        (use string sorting algorithm).
        '''
        if number:
            t = TDBQONUMDESC if desc else TDBQONUMASC
        else:
            t = TDBQOSTRDESC if desc else TDBQOSTRASC
        tctdbqrysetorder(self._state, col, t)
        return self
    
    def limit(self, max=-1, offset=0):
        '''*max* specifies the maximum number of records in the result.
        The default is None.
        
        *offset* specifies the number of skipped records in the result.
        The default is None.
        '''
        if max is None: max = -1
        if offset is None: offset = 0
        tctdbqrysetlimit(self._state, max, offset)
        return self
    
    @property
    def no(self):
        '''Sets up for negation the next condition called on this query.
        Example:
        
            mytable.query.no.endswith('name', 'Palin')
        '''
        self._no = 1
        return self
    
    def _op(self, colname, int flag):
        '''Internal method. Don't call this yourself.'''
        # Determines column name and search flag, consuming instance attributes
        # that may have been previously set on this query object.
        if self._no:
            self._no = 0 # consume _no
            flag = flag | TDBQCNEGATE
        colname = colname or self._col
        self._col = None # consume _col
        return colname, flag
    
    def equals(self, val, col=None):
        '''Adds a condition to the query: the value of column *col* must
        equal *val*. Returns the same query object.
        '''
        col, flag = self._op(col, TDBQCSTREQ)
        tctdbqryaddcond(self._state, col, flag, val)
        return self
    
    def contains(self, val, col=None):
        '''Adds a condition to the query: the value of column *col* must
        contain *val*. Returns the same query object.
        '''
        col, flag = self._op(col, TDBQCSTRINC)
        tctdbqryaddcond(self._state, col, flag, val)
        return self
    
    def startswith(self, val, col=None):
        '''Adds a condition to the query: the value of column *col* must
        begin with *val*. Returns the same query object.
        '''
        col, flag = self._op(col, TDBQCSTRBW)
        tctdbqryaddcond(self._state, col, flag, val)
        return self
    
    def endswith(self, val, col=None):
        '''Adds a condition to the query: the value of column *col* must
        end with *val*. Returns the same query object.
        '''
        col, flag = self._op(col, TDBQCSTREW)
        tctdbqryaddcond(self._state, col, flag, val)
        return self
    
    def contains_all(self, val, col=None):
        '''Adds a condition to the query: the value of column *col* must
        include all tokens in *val*, which may be a list,
        or preferably a string. Returns the same query object.
        '''
        if isinstance(val, list):
            val = ' '.join(val)
        col, flag = self._op(col, TDBQCSTRAND)
        tctdbqryaddcond(self._state, col, flag, val)
        return self
    
    def contains_one(self, val, col=None):
        '''Adds a condition to the query: the value of column *col* must
        include one of the words in *val*, which may be a list,
        or preferably a string.
        
        Returns the same query object.
        '''
        if isinstance(val, list):
            val = ' '.join(val)
        col, flag = self._op(col, TDBQCSTROR)
        tctdbqryaddcond(self._state, col, flag, val)
        return self
    
    def equals_one(self, val, col=None):
        '''Adds a condition to the query: the value of column *col* must
        equal one of the words in *val*, which may be a list,
        or preferably a string.
        
        Returns the same query object.
        '''
        if isinstance(val, list):
            val = ' '.join(val)
        col, flag = self._op(col, TDBQCSTROREQ)
        tctdbqryaddcond(self._state, col, flag, val)
        return self
    
    def matches_regex(self, val, col=None):
        '''Adds a condition to the query: the value of column *col* must
        match the regular expression pattern in *val*.
        Returns the same query object.
        '''
        col, flag = self._op(col, TDBQCSTRRX)
        tctdbqryaddcond(self._state, col, flag, val)
        return self
    
    def number_equals(self, val, col=None):
        '''Adds a condition to the query: the value of column *col* must
        be equal to *val*. Returns the same query object.
        '''
        col, flag = self._op(col, TDBQCNUMEQ)
        tctdbqryaddcond(self._state, col, flag, val)
        return self
    
    def greater_than(self, val, col=None):
        '''Adds a condition to the query: the value of column *col* must
        be greater than the number *val*. Returns the same query object.
        '''
        col, flag = self._op(col, TDBQCNUMGT)
        tctdbqryaddcond(self._state, col, flag, val)
        return self
    
    def greater_or_equal(self, val, col=None):
        '''Adds a condition to the query: the value of column *col* must
        be greater than, or equal to, the number *val*.
        Returns the same query object.
        '''
        col, flag = self._op(col, TDBQCNUMGE)
        tctdbqryaddcond(self._state, col, flag, val)
        return self
    
    def less_than(self, val, col=None):
        '''Adds a condition to the query: the value of column *col* must
        be lower than the number *val*. Returns the same query object.
        '''
        col, flag = self._op(col, TDBQCNUMLT)
        tctdbqryaddcond(self._state, col, flag, val)
        return self
    
    def less_or_equal(self, val, col=None):
        '''Adds a condition to the query: the value of column *col* must
        be lower than, or equal to, the number *val*.
        Returns the same query object.
        '''
        col, flag = self._op(col, TDBQCNUMLE)
        tctdbqryaddcond(self._state, col, flag, val)
        return self
    
    def between(self, low, high, col=None):
        '''Adds a condition to the query: the value of column *col* must
        be between the two numbers *low* and *high*.
        Returns the same query object.
        '''
        val = low + ' ' + high
        col, flag = self._op(col, TDBQCNUMBT)
        tctdbqryaddcond(self._state, col, flag, val)
        return self
    
    def number_in(self, val, col=None):
        '''Adds a condition to the query: the value of column *col* must
        equal one of the numbers in *val*.
        *val* may be a list, or preferrably a string.
        Returns the same query object.
        '''
        if isinstance(val, list):
            val = ' '.join(val)
        col, flag = self._op(col, TDBQCNUMOREQ)
        tctdbqryaddcond(self._state, col, flag, val)
        return self



cdef TCQuery make_query(TCTable table, TCTDB* tctdb):
    # Factory of queries. Before doing it this way, tried passing the above
    # args to the TCQuery constructor, but couldn't make that work.
    # http://www.mail-archive.com/cython-dev@codespeak.net/msg03932.html
    cdef TCQuery query = TCQuery()
    query._table = table
    query._state = tctdbqrynew(tctdb)
    query._col = None
    query._no = False
    return query
